from fastNLP.core.batch import DataSetIter
from fastNLP.core.sampler import RandomSampler, SequentialSampler


class Task(object):
    def __init__(self, task_id, task_name, train_set, dev_set, test_set):
        self.task_id = task_id
        self.task_name = task_name

        self.train_set = train_set
        self.dev_set = dev_set
        self.test_set = test_set

        self.train_loader = None
        self.dev_loader = None
        self.test_loader = None

        self.metric_key = None
        self.metrics = []

    def init_data_loader(self, batch_size):
        self.train_loader = DataSetIter(self.train_set,
                                        batch_size,
                                        sampler=RandomSampler())
        self.dev_loader = DataSetIter(self.dev_set,
                                      batch_size,
                                      sampler=SequentialSampler())
        self.test_loader = DataSetIter(self.test_set,
                                       batch_size,
                                       sampler=SequentialSampler())

    def set_metrics(self, metric_key, metrics):
        self.metric_key = metric_key
        self.metrics = metrics
